#include "aruco_location/aruco_location.hpp"

aruco_location::ArucoLocation::ArucoLocation(const std::string &config_path)
                                                : config_path_(config_path)
{   
    setParamters();
    dict_ = cv::aruco::getPredefinedDictionary(aruco_format_id_);
    result_pose_.reset(new ResultPose);
    pose_array_.reset(new geometry_msgs::PoseArray());

    pub_img_ = n_.advertise<sensor_msgs::Image>(topic_detect_, 10);
    pub_pose_array_ = n_.advertise<geometry_msgs::PoseArray>(topic_pose_, 10);
    
    std::cout << "\033[1;31m" << "-----> waitting image ... " << "\033[0m" << std::endl;
    sub_img_ = n_.subscribe(topic_img_, 10, &ArucoLocation::callback, this);

}

void aruco_location::ArucoLocation::setParamters()
{   
    double fx, fy, cx, cy;
    double k1, k2, p1, p2, k3;

    cv::FileStorage fs(config_path_, cv::FileStorage::READ);
    cv::FileNode n = fs["distortion_parameters"];
    k1             = static_cast< double >( n["k1"] );
    k2             = static_cast< double >( n["k2"] );
    k3             = static_cast< double >( n["k3"] );
    p1             = static_cast< double >( n["p1"] );
    p2             = static_cast< double >( n["p2"] );

    n  = fs["projection_parameters"];
    fx = static_cast< double >( n["fx"] );
    fy = static_cast< double >( n["fy"] );
    cx = static_cast< double >( n["cx"] );
    cy = static_cast< double >( n["cy"] );

    cv_K_ = (cv::Mat_<float>(3, 3) << fx, 0., cx,
				0, fy, cy,
				0, 0., 1);

    cv_Dist_ = (cv::Mat_<float>(5, 1) << k1, k2,
			    p1, p2, k3);

    double x, y, z, roll, pitch, yaw;
    n = fs["Tworld_cam"];
    x     = static_cast< double >( n["x"] );
    y     = static_cast< double >( n["y"] );
    z     = static_cast< double >( n["z"] );
    roll  = static_cast< double >( n["roll"] );
    pitch = static_cast< double >( n["pitch"] );
    yaw   = static_cast< double >( n["yaw"] );

    Eigen::Vector3d eu(yaw, pitch, roll);
    twc_ = Eigen::Vector3d(x, y, z);

    
    eu *= ( M_PI / 180.); 
    euler2Quat(eu, qwc_);

    aruco_len_        = ARUCO_LEN;
    aruco_format_id_  = ARUCO_ID;
    is_draw_pose_     = DRAW_POSE;
    
    fs["aruco_index"] >> aruco_index_;
    fs["topic_img"]    >> topic_img_;
    fs["topic_detect"] >> topic_detect_;
    fs["topic_pose"]   >> topic_pose_;

    fs["world_frame_id"]  >> world_frame_id_;
    fs["camera_frame_id"] >> camera_frame_id_;
    fs["aruco_frame_id"]  >> aruco_frame_id_;
    pub_TF_flag_ = static_cast<int>(fs["pub_TF_flag"]);
    pub_image_flage_  = static_cast<int>(fs["pub_image_flage"]);

    std::cout << std::fixed << std::setprecision(2) 
                << "\033[1;32m\n----->world2cam(rpyxyz):" 
                << roll << "  " << pitch << "  " << yaw << "  " 
                << twc_(0) << "  " << twc_(1) << "  " << twc_(2) << "\033[0m" << std::endl;

     std::cout << std::fixed << std::setprecision(2) << "\033[1;32m-----> fx fy cx cy     :" 
                << fx << "  "
                << fy << "  "
                << cx << "  "
                << cy << "  "
                << "\033[0m" << std::endl;
    
    std::cout << std::fixed << std::setprecision(2) << "\033[1;32m-----> k1 k2 p1 p2 k3  :" 
                << k1 << "  "
                << k2 << "  "
                << p1 << "  "
                << p2 << "  "
                << k3 << "  "
                << "\033[0m" << std::endl;
    
    std::cout << std::fixed << std::setprecision(3) << "\n\033[1;32m-----> aruco_length: " << ARUCO_LEN << "m\033[0m" << std::endl;
    std::cout << "\033[1;32m-----> aruco_id    : " << ARUCO_ID << "\033[0m" << std::endl;
    std::cout << "\033[1;32m-----> aruco_index : " << aruco_index_ << "\033[0m" << std::endl;
    std::cout << "\033[1;32m-----> topic_img   : " << topic_img_ << "\033[0m" << std::endl;
    std::cout << "\033[1;32m-----> topic_detect: " << topic_detect_ << "\033[0m" << std::endl;
    std::cout << "\033[1;32m-----> topic_pose  : " << topic_pose_ << "\n\033[0m" << std::endl;
    
    std::cout << "\033[1;32m-----> world_frame_id : " << world_frame_id_ << "\033[0m" << std::endl;
    std::cout << "\033[1;32m-----> camera_frame_id: " << camera_frame_id_ << "\n\033[0m" << std::endl;
    std::cout << "\033[1;32m-----> aruco_frame_id : " << aruco_frame_id_ << "\033[0m" << std::endl;
    std::cout << "\033[1;32m-----> draw_pose      : " << DRAW_POSE << "\033[0m" << std::endl;
    std::cout << "\033[1;32m-----> pub_TF_flag    : " << pub_TF_flag_ << "\033[0m" << std::endl;
    std::cout << "\033[1;32m-----> pub_image_flage: " << pub_image_flage_ << "\n\033[0m" << std::endl;
}

bool aruco_location::ArucoLocation::detectAruco()
{
    cv::Mat img = img_.clone();
    draw_img_ = img_.clone();
    
    // 1 检测Marker标志，返回id和corner的列表
    // cv::aruco::detectMarkers(inputImage, dict, markerCorners, markerIds, parameters, rejectedCandidates);
    cv::aruco::detectMarkers(img, dict_, marker_corners_, markerIds_);
    
    if(markerIds_.size()>0)
    {   
        // 3 估计位姿, 第二个参数markerLength边长,单位m
        cv::aruco::estimatePoseSingleMarkers(marker_corners_, aruco_len_, cv_K_, cv_Dist_, result_pose_->rvecs, result_pose_->tvecs);
        
        
        // 4 绘制坐标轴，检查姿态估计结果
        geometry_msgs::Pose pose;
        
        for (int i = 0; i < result_pose_->rvecs.size(); i++)
        { 
            cv::Mat R_cv = cv::Mat::eye(3, 3, CV_64F);
            cv::Rodrigues(result_pose_->rvecs[i], R_cv);
            
            Eigen::Matrix3d R_eigen = Eigen::Matrix3d::Identity();
            cv::cv2eigen(R_cv, R_eigen);
            Eigen::Quaterniond q_ca(R_eigen);
            
            Eigen::Vector3d tca(result_pose_->tvecs[i](0), result_pose_->tvecs[i](1), result_pose_->tvecs[i](2));
            Eigen::Quaterniond qwa = qwc_ * q_ca;
            qwa.normalize();
            Eigen::Vector3d twa = qwc_ * tca + twc_;

            pose_array_->header.frame_id = aruco_frame_id_;
            pose_array_->header.stamp = ros::Time::now();
        
            pose.position.x = twa.x();
            pose.position.y = twa.y();
            pose.position.z = twa.z();

            tf2::Quaternion _q(qwa.x(), qwa.y(), qwa.z(), qwa.w());
            // tf2::Quaternion _q(qca.x(), qca.y(), qca.z(), qca.w());
            double _r, _p, _y;
            tf2::Matrix3x3(_q).getRPY(_r, _p, _y);

            pose.orientation.x = _r;
            pose.orientation.y = _p;
            pose.orientation.z = _y;
            pose.orientation.w = markerIds_[i];
            pose_array_->poses.emplace_back(pose);
            
            // 可视化
            if(is_draw_pose_)
            {
                // 1 这个只能画框
                cv::aruco::drawDetectedMarkers(img, marker_corners_, markerIds_);

                // 2 绘制坐标轴，检查姿态估计结果 // 最后一个参数为坐标轴长度单位m
                cv::aruco::drawAxis(img, cv_K_, cv_Dist_, result_pose_->rvecs[i], result_pose_->tvecs[i], 0.1);         
                    
                // // 3 图像上画圆心
                // cv::Point2i cricle_center( (marker_corners_[i][0].x + marker_corners_[i][2].x)/2,
                //                         (marker_corners_[i][0].y + marker_corners_[i][2].y)/2 );
                    
                // cv::circle(img, cricle_center, 4, cv::Scalar(0, 255, 255), -1, 4 );
            }
            img.copyTo(draw_img_);
            // draw_img_ = img.clone();
        }
        return true;
    }
    else 
    {   
         std::cout << "\033[1;31m" << "-----> Detection failure ... " << "\033[0m" << std::endl;
        return false;  // 未检测到
    }
    
}

void aruco_location::ArucoLocation::callback(const sensor_msgs::ImageConstPtr &msg)
{   
    

    // 检测的二维码
    img_ = cv_bridge::toCvShare(msg, "bgr8")->image;
    cv::equalizeHist(img_, img_);
    // cv::cvtColor(img_, img_, cv::COLOR_BGR2GRAY);
    // std::cout << img_.size() << std::endl;
     
    // 检测 会更新 markerIds_ marker_corners_, rejected_candidates_, result_pose_
    pose_array_->poses.clear();
    bool ret = detectAruco();
    
    double s = 180. / M_PI;
    
    if(ret)
       
        pub_pose_array_.publish(pose_array_);
    
    if(ret && pub_TF_flag_)
    {

        // 发布TF
        geometry_msgs::TransformStamped tf_wc;
        tf_wc.child_frame_id = camera_frame_id_;
        tf_wc.header.frame_id = world_frame_id_;
        tf_wc.header.stamp = ros::Time::now();
        tf_wc.transform.translation.x = twc_.x();
        tf_wc.transform.translation.y = twc_.y();
        tf_wc.transform.translation.z = twc_.z();
        tf_wc.transform.rotation.w = qwc_.w();
        tf_wc.transform.rotation.x = qwc_.x();
        tf_wc.transform.rotation.y = qwc_.y(); 
        tf_wc.transform.rotation.z = qwc_.z(); 
        broadcaster_.sendTransform(tf_wc);

        for(auto it : pose_array_->poses)
        {   
            tf2::Quaternion q;
            q.setRPY(it.orientation.x, it.orientation.y, it.orientation.z);
            q.normalized();

            std::cout << "markerID: " << int(it.orientation.w) << " " 
                << it.position.x << " " << it.position.y << " " << it.position.z << " "
                << it.orientation.x * s  << " " << it.orientation.y * s << " " << it.orientation.z * s
                << std::endl;
            
            // 发布index为aruco_index_的姿态
            if(it.orientation.w == aruco_index_)
            {
                geometry_msgs::TransformStamped tf_1;
                tf_1.child_frame_id = aruco_frame_id_;
                tf_1.header.frame_id = world_frame_id_;
                tf_1.header.stamp = ros::Time::now();
                tf_1.transform.translation.x = it.position.x;
                tf_1.transform.translation.y = it.position.y;
                tf_1.transform.translation.z = it.position.z;
                tf_1.transform.rotation.w = q.w();
                tf_1.transform.rotation.x = q.x();
                tf_1.transform.rotation.y = q.y(); 
                tf_1.transform.rotation.z = q.z(); 
                broadcaster_.sendTransform(tf_1);
            }
        }
    }
    std::cout << std::endl; 
    
    // 发布带坐标的检测图像,没检测到就发原图
    sensor_msgs::ImagePtr msg_img;
    if(!ret)
        cv::putText(draw_img_,"No ARUCO detected.", cv::Point2i(20,50), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 255,0), 2, 4);
    msg_img = cv_bridge::CvImage(std_msgs::Header(), "bgr8", draw_img_).toImageMsg();
    
    if(pub_image_flage_)
        pub_img_.publish(msg_img);
}